'''
2-1 构建样本集
样本数 m = 100
按照7:3的比例创建训练集(70)和测试集(30)，训练集(0-69),前半部分(0-34)放猫的图像，后半部分(35-69)放狗的图像，测试集(0-29),前半部分(0-14)放猫的图像，后半部分(15-29)放狗的图像
cv.resize() 对图像进行缩放到规定尺寸
numpy 创建零矩阵，存放训练集和测试集
os模块 判断文件路径
'''
import cv2 as cv
import numpy as np
import os
if __name__ == '__main__':
    l_x = 128    # 输入x长宽
    m = 100      # 样本数量
    m1 = 35
    m2 = 15
    train_set = np.zeros(m1*2*l_x*l_x*3)
    train_set = np.reshape(train_set,(m1*2,3,l_x,l_x))
    test_set = np.zeros(m2*2*l_x*l_x*3)
    test_set = np.reshape(test_set,(m2*2,3,l_x,l_x))  #创建零矩阵，存放训练集和测试集
    sucess_mark = 0          #成功标记

    #构建训练集
    for i in range(m1):          # m1是0到35，循环一次猫图像被放在第 i 个位置上，狗的图像被放在第 m1+i 个位置上
        path1 = f'./sample/cat.{i}.jpg'    # 路径
        path2 = f'./sample/dog.{i}.jpg'
        if os.path.exists(path1) & os.path.exists(path2):    #判断路径是否存在
            img1 = cv.imread(path1)    # 读取一张猫的图片
            img2 = cv.imread(path2)    # 读取一张狗的图片
            img1 = cv.resize(img1,(l_x,l_x))
            img2 = cv.resize(img2,(l_x,l_x))  #将图像缩放到规定尺寸
            train_set[i, 0, :, :] = img1[:, :, 0]
            train_set[i, 1, :, :] = img1[:, :, 1]
            train_set[i, 2, :, :] = img1[:, :, 2]  #赋值给训练集矩阵，三通道, 猫
            sucess_mark += 1
            print("\r" + f'训练集总数:{m1*2},当前第{(i+1)*2}个',end = "",flush = True)
            train_set[m1+i, 0, :, :] = img2[:, :, 0]
            train_set[m1+i, 1, :, :] = img2[:, :, 1]
            train_set[m1+i, 2, :, :] = img2[:, :, 2]  # 赋值给训练集矩阵，三通道， 狗
            sucess_mark += 1
        else:
            print(f'路径{path1}或{path2}不存在')
            break
    print('') #换行

    #构建测试集
    for i in range(35,50):          #第一次循环猫的第35张图像在测试集m2的第0个位置上，狗的第35张图像在测试集m2的第15个位置上
        path1 = f'./sample/cat.{i}.jpg'    # 路径
        path2 = f'./sample/dog.{i}.jpg'
        if os.path.exists(path1) & os.path.exists(path2):    #判断路径是否存在
            img1 = cv.imread(path1)    # 读取一张猫的图片
            img2 = cv.imread(path2)    # 读取一张狗的图片
            img1 = cv.resize(img1,(l_x,l_x))
            img2 = cv.resize(img2,(l_x,l_x))  #将图像缩放到规定尺寸
            test_set[i-35, 0, :, :] = img1[:, :, 0]
            test_set[i-35, 1, :, :] = img1[:, :, 1]
            test_set[i-35, 2, :, :] = img1[:, :, 2]  #赋值给训练集矩阵，三通道, 猫
            sucess_mark += 1
            print("\r" + f'测试集总数:{(i-35+1)*2},当前第{(i+1)*2}个',end = "",flush = True)
            test_set[m2+i-35, 0, :, :] = img2[:, :, 0]
            test_set[m2+i-35, 1, :, :] = img2[:, :, 1]
            test_set[m2+i-35, 2, :, :] = img2[:, :, 2]  # 赋值给训练集矩阵，三通道， 狗
            sucess_mark += 1
        else:
            print(f'路径{path1}或{path2}不存在')
            break
    print('') #换行
    if sucess_mark == 100:    #如果标记个数等于样本总数，则保存训练集和测试集
        np.save('cat_train_set',train_set)
        np.save('cat_test_set',test_set)
        print('生成成功')